import argparse
import os
import sys
from pathlib import Path

import torchvision
import torchvision.transforms as transforms
from torch.optim import SGD
from torch.utils.data import DataLoader
from mmengine.runner import Runner

from resnet.model import MMResNet50
from resnet.metric import Accuracy


FILE = Path(__file__).resolve()
ROOT = FILE.parents[0]  # root directory
if str(ROOT) not in sys.path:
    sys.path.append(str(ROOT))  # add ROOT to PATH
ROOT = Path(os.path.relpath(ROOT, Path.cwd()))  # relative


def main(opt):
    # define the dataloaders
    norm_cfg = dict(mean=[0.491, 0.482, 0.447], std=[0.202, 0.199, 0.201])
    train_dataloader = DataLoader(
        batch_size=32,
        shuffle=True,
        dataset=torchvision.datasets.CIFAR10(
            opt.data_dir,
            train=True,
            download=True,
            transform=transforms.Compose(
                [
                    transforms.RandomCrop(32, padding=4),
                    transforms.RandomHorizontalFlip(),
                    transforms.ToTensor(),
                    transforms.Normalize(**norm_cfg),
                ]
            ),
        ),
    )

    val_dataloader = DataLoader(
        batch_size=32,
        shuffle=False,
        dataset=torchvision.datasets.CIFAR10(
            opt.data_dir,
            train=False,
            download=True,
            transform=transforms.Compose(
                [transforms.ToTensor(), transforms.Normalize(**norm_cfg)]
            ),
        ),
    )

    # define the runner and train
    runner = Runner(
        model=MMResNet50(),
        work_dir=opt.work_dir,
        train_dataloader=train_dataloader,
        optim_wrapper=dict(optimizer=dict(type=SGD, lr=0.001, momentum=0.9)),
        train_cfg=dict(by_epoch=True, max_epochs=opt.epochs, val_interval=1),
        val_dataloader=val_dataloader,
        val_cfg=dict(),
        val_evaluator=dict(type=Accuracy),
    )
    runner.train()


def parse_opt():
    parser = argparse.ArgumentParser()
    parser.add_argument(
        "--data_dir", default=ROOT / "data/cifar10", help="the directory of data"
    )
    parser.add_argument(
        "--work_dir", default=ROOT / "runs/resnet/", help="the directory of running"
    )
    parser.add_argument("--epochs", default=5, type=int, help="epochs to train")
    opt = parser.parse_args()
    return opt


if __name__ == "__main__":
    opt = parse_opt()
    main(opt)
